package org.gatech.linalg;

import java.io.*;
import java.util.StringTokenizer;
import java.util.ArrayList;

/**
 * Image compression and decompression codec using singular value decomposition
 * and k-rank approximation. Can be handed any ppm or pgm files and will create
 * compressed versions of the files, with extensions cppm or cpgm respectively
 * Conversely, can be handed a compressed cppm or cpgm file, and will create the
 * decompressed images, with ppm and pgm extensions and with "_deflated"
 * appended to the end of the file names.
 * 
 * @author Taylor Wrobel, Daniel Whatley, Stephen Roca
 * 
 */
public class ImageCompression {

	public static final int ENCODE_G = 0, ENCODE_RGB = 1, DECODE_G = 2,
			DECODE_RGB = 3;
	public static final int READ_TYPE = 0, READ_WIDTH = 1, READ_HEIGHT = 2,
			READ_MAX_VALUE = 3, READ_DATA = 4;

	public static boolean mFlag, eFlag, aFlag, isRGB;
	public static int kRank, iVal;
	public static String filename, gFileName, bFileName;

	/**
	 * Main thread of execution for the codec.
	 * 
	 * @param args
	 *            Command line arguments array
	 */
	public static void main(String[] args) {

		/*
		 * Process command line arguments and set condition variables as needed;
		 */

		if (args.length == 0)
			showHelp();

		filename = args[0];
		kRank = -1;
		iVal = -1;
		isRGB = false;
		mFlag = false;
		eFlag = false;
		aFlag = false;
		for (int i = 1; i < args.length; i++) {
			if (args[i].equals("-a"))
				aFlag = true;
			else if (args[i].equals("-e"))
				eFlag = true;
			else if (args[i].equals("-h"))
				showHelp();

			else if (args[i].equals("-i")) {
				try {
					iVal = Integer.parseInt(args[i + 1]);
					i++;
				} catch (Exception e) {
					showHelp();
				}
			}

			else if (args[i].equals("-irgb")) {
				try {
					iVal = Integer.parseInt(args[i + 1]);
					gFileName = args[i + 2];
					bFileName = args[i + 3];
					isRGB = true;
					i += 3;
				} catch (Exception e) {
					showHelp();
				}
			}

			else if (args[i].equals("-m"))
				mFlag = true;
			else if (args[i].equals("-k")) {
				try {
					kRank = Integer.parseInt(args[i + 1]);
					i++;
				} catch (Exception e) {
					showHelp();
				}
			} else
				showHelp();
		}

		if (!eFlag && iVal == -1)
			transcode();
		if (iVal != -1)
			writeImage();
	}

	/**
	 * Displays the help dialog with all the flag options and then exits the
	 * system.
	 */
	public static void showHelp() {
		System.err.println("Usage and Flags:");
		System.err.println("ImageCompression \"filename\" [-FLAGS]");
		System.err.println("\tflags:");
		System.err
				.println("\t-a Convert the image with name \"filename\" to a k-rank"
						+ "\n\t\tapproximation of the image.  If k is not specified with the"
						+ "\n\t\t-k flag it will be automatically determined.");
		System.err
				.println("\t-e Output the eigenvalues and eigenvectors of the matrix file "
						+ "\n\t\t" + "with name \"filename\"");
		System.err.println("\t-h Display this help dialog");
		System.err
				.println("\t-i [int] Output the matrix file with name \"filename\" to an image "
						+ "\n\t\t" + " file with maximum value [int]");
		System.err
				.println("\t-irgb [int] \"filename2\" \"filename3\" Outputs an RGB image, using the"
						+ "\n\t\t"
						+ "matrix file with name \"filename\" as the red channel, the one "
						+ "\n\t\t"
						+ "with name \"filename2\" as the green channel, and the one with "
						+ "\n\t\t"
						+ "name \"filename3\"as the blue channel, and with maximum color "
						+ "\n\t\t" + "value [int]");
		System.err
				.println("\t-k [int] Set the k to be used in k-rank approximations");
		System.err
				.println("\t-m Output the ppm or pgm image with name \"filename\" to a matrix file");

		System.exit(0);
	}

	/**
	 * Support for -i flag
	 */
	public static void writeImage() {

		int count = 0;
		String[] filenames = null;
		Matrix[] matrices = new Matrix[3];
		if (isRGB) {
			count = 3;
			String[] files = { filename, gFileName, bFileName };
			filenames = files;
		} else {
			count = 1;
			String[] files = { filename };
			filenames = files;
		}

		for (int loopNum = 0; loopNum < count; loopNum++) {

			BufferedReader in = null;
			try {
				in = new BufferedReader(new FileReader(new File(
						filenames[loopNum])));
				System.err.println("Opening file " + filenames[loopNum]);
			} catch (FileNotFoundException e) {
				System.err.println("File \"" + filename + "\" not found.");
				System.exit(1);
			}

			ArrayList<ArrayList<Double>> valueLists = new ArrayList<ArrayList<Double>>();
			try {
				String line = in.readLine();

				while (line != null) {
					ArrayList<Double> row = new ArrayList<Double>();
					StringTokenizer st = new StringTokenizer(line, " ");
					while (st.hasMoreTokens()) {
						String number = st.nextToken();
						while (number.equals(""))
							number = st.nextToken();
						try {
							row.add(Double.parseDouble(number));
						} catch (NumberFormatException e) {
							System.err
									.println("Malformed matrix file.  Exiting.");
							System.exit(42);
						}
					}
					valueLists.add(row);
					line = in.readLine();
				}
			} catch (IOException e) {
				System.err.println("Problem reading from file " + filename);
				System.exit(42);
			}

			double[][] valueArray = new double[valueLists.size()][valueLists
					.get(0).size()];
			for (int i = 0; i < valueArray.length; i++) {
				for (int j = 0; j < valueArray[0].length; j++) {
					valueArray[i][j] = valueLists.get(i).get(j);
				}
			}

			try {
				matrices[loopNum] = new Matrix(valueArray);
			} catch (DimensionMismatchException e) {
				System.err
						.println("Dimensionality error with one or more matrices.");
				System.exit(3);
			}
		}

		try {
			String outputName = filename
					.substring(0, filename.lastIndexOf('.'))
					+ (isRGB ? ".ppm" : ".pgm");
			if (!isRGB) {
				if (eFlag) {
					calcEigsystem(matrices[0]);
					System.exit(0);
				}
				MatrixWriter.writeGrayscale(matrices[0], outputName, iVal);
			} else
				MatrixWriter.writeRGB(matrices, outputName, iVal);
		} catch (Exception e) {
			System.err.println("Error when trying to write image file.");
			e.printStackTrace();
		}
	}

	public static void transcode() {
		/*
		 * Read in file and preprocess, removing all comments, tho can still
		 * have redundant white space.
		 */
		BufferedReader in = null;
		try {
			in = new BufferedReader(new FileReader(new File(filename)));
		} catch (FileNotFoundException e) {
			System.err.println("File \"" + filename + "\" not found.");
			System.exit(1);
		}

		String[] uncommented = preProcess(in);

		/*
		 * Set up instance and state variables, as well as data for the arrays
		 * and matrices
		 */
		int opType = ENCODE_G;
		int dataType = READ_TYPE;
		int width = 0, height = 0;
		int row = 0, col = 0, maxVal = 0, rgb = 0;
		double[][] twoArr = null;
		// double[][][][] megaArray = new double[3][3][][];
		double[][][] threeArr = null;

		for (int line = 0; line < uncommented.length; line++) {

			/*
			 * Establish String Tokenizer for parsing file information.
			 */
			StringTokenizer st = new StringTokenizer(uncommented[line]);

			/*
			 * Parse through file in terms of tokens
			 */
			while (st.hasMoreTokens()) {
				String token = st.nextToken();
				while (token.equals(""))
					token = st.nextToken();

				/*
				 * Switch on the current type of data being parsed, based on
				 * states Operations done on data can also switch on operation
				 * type, after it has been determined.
				 */
				switch (dataType) {
				case READ_TYPE:
					if (token.equalsIgnoreCase("p2"))
						opType = ENCODE_G;
					else if (token.equalsIgnoreCase("p3"))
						opType = ENCODE_RGB;
					else if (token.equalsIgnoreCase("cp2"))
						opType = DECODE_G;
					else if (token.equalsIgnoreCase("cp3"))
						opType = DECODE_RGB;
					else {
						System.err
								.println("Not a correct filetype for the operation selected.");
						System.exit(0);
					}
					dataType = READ_WIDTH;
					break;
				case READ_WIDTH:
					width = Integer.parseInt(token);
					dataType = READ_HEIGHT;
					break;
				case READ_HEIGHT:
					height = Integer.parseInt(token);
					dataType = READ_MAX_VALUE;
					switch (opType) {
					case ENCODE_G:
						twoArr = new double[height][width];
						break;
					case ENCODE_RGB:
						threeArr = new double[3][height][width];
						break;
					case DECODE_G:
						twoArr = new double[height][width];
						break;
					case DECODE_RGB:
						threeArr = new double[3][height][width];
						break;
					}
					break;
				case READ_MAX_VALUE:
					maxVal = Integer.parseInt(token);
					dataType = READ_DATA;
					break;
				case READ_DATA:
					switch (opType) {
					case ENCODE_G:
						twoArr[row][col] = Integer.parseInt(token);
						col++;
						row += col / width;
						col = col % width;
						break;
					case ENCODE_RGB:
						threeArr[rgb][row][col] = Integer.parseInt(token);
						rgb++;
						if (rgb == 3) {
							col++;
							row += col / width;
							col = col % width;
						}
						rgb %= 3;
						break;
					}
					break;
				}
			}// End while
		}

		/*
		 * Done parsing file. Now opType can be used to manipulate pertinent
		 * data and perform proper operations.
		 */

		switch (opType) {
		case ENCODE_G:
			try {
				Matrix mat = new Matrix(twoArr);
				if (kRank == -1)
					kRank = generateKRank(mat);
				if (mFlag) {
					Matrix[] toFiles = { mat };
					makeMatrixFiles(toFiles);
				}
				// Take care of handling the -a flag
				else if (aFlag) {
					System.err.println("DOING A " + kRank
							+ "-RANK APPROXIMATION");
					Matrix approximated = kRankApprox(mat);

					String outputName = filename.substring(0,
							filename.lastIndexOf('.'))
							+ "_approximated.pgm";
					try {
						MatrixWriter.writeGrayscale(approximated, outputName,
								maxVal);
					} catch (IOException e) {
						System.err.println("Error writing to file "
								+ outputName);
					}
				}
			} catch (DimensionMismatchException e) {
				e.printStackTrace();
			}
			break;
		case ENCODE_RGB:
			try {
				Matrix red = new Matrix(threeArr[0]);
				Matrix green = new Matrix(threeArr[1]);
				Matrix blue = new Matrix(threeArr[2]);
				if (kRank == -1)
					kRank = generateKRank(red);
				if (mFlag) {
					System.err.println("Test 3");
					Matrix[] toFiles = { red, green, blue };
					makeMatrixFiles(toFiles);
				}
				// Take care of handling the -a flag
				else if (aFlag) {
					System.err.println("DOING A " + kRank
							+ "-RANK APPROXIMATION ON AN RGB IMAGE");
					Matrix approximatedRed = kRankApprox(red);
					Matrix approximatedGreen = kRankApprox(green);
					Matrix approximatedBlue = kRankApprox(blue);

					String outputName = filename.substring(0, (filename
							.indexOf('_') == -1 ? filename.lastIndexOf('.')
							: filename.lastIndexOf('_')))
							+ "_approximated.ppm";
					Matrix[] approximated = { approximatedRed,
							approximatedGreen, approximatedBlue };
					try {
						MatrixWriter.writeRGB(approximated, outputName, maxVal);
					} catch (IOException e) {
						System.err.println("Error writing to file "
								+ outputName);
					}
				}
			} catch (DimensionMismatchException e) {
				e.printStackTrace();
			}
		case DECODE_G:
			// try{
			// Matrix U = new Matrix(megaArray[0][0]);
			// Matrix D = new Matrix(megaArray[0][1]);
			// Matrix V = new Matrix(megaArray[0][2]);
			//
			// System.out.println(V);
			//
			// Matrix[] partials = {U,D,V};
			//
			// Matrix remade = remakeKApprox(partials);
			//
			// String outputName = filename.substring(0, (filename
			// .indexOf('_') == -1 ? filename.lastIndexOf('.')
			// : filename.lastIndexOf('_')))
			// + "_deflated.pgm";
			// MatrixWriter.writeGrayscale(remade, outputName, maxVal);
			// } catch(DimensionMismatchException e){
			// System.err.println("WRONG DIMENSIONS, FOOL!");
			// } catch(IOException e){
			// System.err.println("IO ERROR?!");
			// }
			// break;
		case DECODE_RGB:

			break;
		}
	}

	/*
	 * Writes an array of given matrices to .mat files. If only one matrix is
	 * given, will only write one file. If three matrices are given, will write
	 */
	public static void makeMatrixFiles(Matrix[] matrices) {
		if (!eFlag) {
			for (int i = 0; i < matrices.length; i++) {
				String outputName = filename.substring(0,
						filename.lastIndexOf('.'));
				if (i == 0)
					if (matrices.length == 1)
						outputName += ".mat";
					else
						outputName += "_red.mat";
				else if (i == 1)
					outputName += "_green.mat";
				else if (i == 2)
					outputName += "_blue.mat";
				else {
					System.exit(1337);
				}
				try {
					PrintWriter out = new PrintWriter(new File(outputName));

					double[][] data = matrices[i].getVals();

					for (int j = 0; j < data.length; j++) {
						for (int k = 0; k < data[j].length; k++) {
							out.print(data[j][k]
									+ (k == data[j].length - 1 ? "" : " "));
						}
						out.print((j == data.length - 1 ? "" : "\n"));
					}

					System.err.println("Wrote matrix file with dimensions "
							+ data.length + ", " + data[0].length);

					out.flush();
					out.close();

				} catch (IOException e) {
					System.err.println("Error writing to file " + outputName);
				}
			}
		}
	}

	private static Matrix kRankApprox(Matrix m) {
		if (kRank == -1)
			kRank = generateKRank(m);
		if (kRank > Math.min(m.getRows(), m.getCols()))
			kRank = Math.min(m.getRows(), m.getCols());
		double[][] runningVals = new double[m.getRows()][m.getCols()];
		try {
			System.err.println("Starting SVD computation");
			Matrix[] svd = Matrix.SVD(m);
			System.err.println("SVD done computing");
			double[][] U = svd[0].getVals();
			double[][] S = svd[1].getVals();
			double[][] V = svd[2].getVals();

			System.err.println("U is " + U.length + " by " + U[0].length);
			System.err.println("S is " + S.length + " by " + S[0].length);
			System.err.println("V is " + V.length + " by " + V[0].length);

			// System.out.println(Matrix.product(Matrix.product(svd[0], svd[1]),
			// Matrix.transpose(svd[2])));
			for (int i = 0; i < kRank; i++) {
				double[][] Ucol = new double[U.length][1];
				for (int j = 0; j < U.length; j++) {
					Ucol[j][0] = U[j][i];
				}
				double[][] Vcol = new double[V.length][1];
				for (int j = 0; j < V.length; j++) {
					Vcol[j][0] = V[j][i];
				}
				double eig = S[i][i];
				Matrix Umat = new Matrix(Ucol);
				Matrix Vmat = new Matrix(Vcol);
				Matrix iteration = Matrix.scalarMult(
						Matrix.product(Umat, Matrix.transpose(Vmat)), eig);
				for (int j = 0; j < iteration.getRows(); j++)
					for (int k = 0; k < iteration.getCols(); k++)
						runningVals[j][k] += iteration.getVals()[j][k];
				// System.err.println("Loop 1 of " + kRank + " complete.");
			}
			System.out.println("Sigma 1 = " + S[0][0]);
			try {
				System.out.println("Sigma K+1 = " + S[kRank + 1][kRank + 1]);
			} catch (ArrayIndexOutOfBoundsException e) {
				System.out.println("Sigma K+1 = 0");
			}
			return new Matrix(runningVals);
		} catch (DimensionMismatchException e) {
			e.printStackTrace();
			System.exit(76);
		}
		return null;
	}

	/**
	 * Pre-processes the file, removing comments and new lines. Will still leave
	 * redundancies in white space!
	 * 
	 * @param file
	 *            The BufferedReader for the file being processed.
	 * @return A String of the file, all on one line, and with no comments.
	 */
	private static String[] preProcess(BufferedReader file) {
		String parsed = "";
		ArrayList<String> lines = new ArrayList<String>();
		try {
			String line = file.readLine();
			while (line != null && !line.equals("")) {
				for (int i = 0; i < line.length(); i++) {
					if (line.charAt(i) == '#')
						break;
					else
						parsed += "" + line.charAt(i);
				}
				lines.add(parsed);
				parsed = "";
				line = file.readLine();
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
		Object[] strObjects = lines.toArray();
		String[] lineArr = new String[strObjects.length];
		for (int i = 0; i < strObjects.length; i++)
			lineArr[i] = (String) strObjects[i];
		return lineArr;
	}

	private static int generateKRank(Matrix m) {
		return (Math.min(m.getRows(), m.getCols()) + 1) / 4;
	}

	public static void calcEigsystem(Matrix m) {
		try{
		Matrix[] jacobiResults = Matrix.QRIteration(m, 1e-4);
		makeEigsystemMatrices(jacobiResults);
		} catch(Exception e){
			e.printStackTrace();
		}
	}
}
